// Copyright (c) 2012 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "media/filters/h264_to_annex_b_bitstream_converter.h"

#include <stddef.h>

#include "base/logging.h"
#include "media/filters/h264_parser.h"
#include "media/formats/mp4/box_definitions.h"

namespace media {

static const uint8_t kStartCodePrefix[3] = { 0, 0, 1 };
static const uint32_t kParamSetStartCodeSize = 1 + sizeof(kStartCodePrefix);

// Helper function which determines whether NAL unit of given type marks
// access unit boundary.
static bool IsAccessUnitBoundaryNal(int nal_unit_type)
{
    // Check if this packet marks access unit boundary by checking the
    // packet type.
    if (nal_unit_type == 6 || // Supplemental enhancement information
        nal_unit_type == 7 || // Picture parameter set
        nal_unit_type == 8 || // Sequence parameter set
        nal_unit_type == 9 || // Access unit delimiter
        (nal_unit_type >= 14 && nal_unit_type <= 18)) { // Reserved types
        return true;
    }
    return false;
}

H264ToAnnexBBitstreamConverter::H264ToAnnexBBitstreamConverter()
    : configuration_processed_(false)
    , first_nal_unit_in_access_unit_(true)
    , nal_unit_length_field_width_(0)
{
}

H264ToAnnexBBitstreamConverter::~H264ToAnnexBBitstreamConverter() { }

bool H264ToAnnexBBitstreamConverter::ParseConfiguration(
    const uint8_t* configuration_record,
    int configuration_record_size,
    mp4::AVCDecoderConfigurationRecord* avc_config)
{
    DCHECK(configuration_record);
    DCHECK_GT(configuration_record_size, 0);
    DCHECK(avc_config);

    if (!avc_config->Parse(configuration_record, configuration_record_size))
        return false; // Error: invalid input

    // We're done processing the AVCDecoderConfigurationRecord,
    // store the needed information for parsing actual payload
    nal_unit_length_field_width_ = avc_config->length_size;
    configuration_processed_ = true;
    return true;
}

uint32_t H264ToAnnexBBitstreamConverter::GetConfigSize(
    const mp4::AVCDecoderConfigurationRecord& avc_config) const
{
    uint32_t config_size = 0;

    for (size_t i = 0; i < avc_config.sps_list.size(); ++i)
        config_size += kParamSetStartCodeSize + avc_config.sps_list[i].size();

    for (size_t i = 0; i < avc_config.pps_list.size(); ++i)
        config_size += kParamSetStartCodeSize + avc_config.pps_list[i].size();

    return config_size;
}

uint32_t H264ToAnnexBBitstreamConverter::CalculateNeededOutputBufferSize(
    const uint8_t* input,
    uint32_t input_size,
    const mp4::AVCDecoderConfigurationRecord* avc_config) const
{
    uint32_t output_size = 0;
    uint32_t data_left = input_size;
    bool first_nal_in_this_access_unit = first_nal_unit_in_access_unit_;

    if (input_size == 0)
        return 0; // Error: invalid input data

    if (!configuration_processed_) {
        return 0; // Error: configuration not handled, we don't know nal unit width
    }

    if (avc_config)
        output_size += GetConfigSize(*avc_config);

    CHECK(nal_unit_length_field_width_ == 1 || nal_unit_length_field_width_ == 2 || nal_unit_length_field_width_ == 4);

    // Then add the needed size for the actual packet
    while (data_left > 0) {
        if (data_left < nal_unit_length_field_width_) {
            return 0; // Error: not enough data for correct conversion.
        }

        // Read the next NAL unit length from the input buffer
        uint8_t size_of_len_field;
        uint32_t nal_unit_length;
        for (nal_unit_length = 0, size_of_len_field = nal_unit_length_field_width_;
             size_of_len_field > 0;
             input++, size_of_len_field--, data_left--) {
            nal_unit_length <<= 8;
            nal_unit_length |= *input;
        }

        if (nal_unit_length == 0) {
            break; // Signifies that no more data left in the buffer
        } else if (nal_unit_length > data_left) {
            return 0; // Error: Not enough data for correct conversion
        }
        data_left -= nal_unit_length;

        // five least significant bits of first NAL unit byte signify nal_unit_type
        int nal_unit_type = *input & 0x1F;
        if (first_nal_in_this_access_unit || IsAccessUnitBoundaryNal(nal_unit_type)) {
            output_size += 1; // Extra zero_byte for these nal units
            first_nal_in_this_access_unit = false;
        }
        // Start code prefix
        output_size += sizeof(kStartCodePrefix);
        // Actual NAL unit size
        output_size += nal_unit_length;
        input += nal_unit_length;
        // No need for trailing zero bits
    }
    return output_size;
}

bool H264ToAnnexBBitstreamConverter::ConvertAVCDecoderConfigToByteStream(
    const mp4::AVCDecoderConfigurationRecord& avc_config,
    uint8_t* output,
    uint32_t* output_size)
{
    uint8_t* out = output;
    uint32_t out_size = *output_size;
    *output_size = 0;
    for (size_t i = 0; i < avc_config.sps_list.size(); ++i) {
        if (!WriteParamSet(avc_config.sps_list[i], &out, &out_size))
            return false;
    }

    for (size_t i = 0; i < avc_config.pps_list.size(); ++i) {
        if (!WriteParamSet(avc_config.pps_list[i], &out, &out_size))
            return false;
    }

    nal_unit_length_field_width_ = avc_config.length_size;
    configuration_processed_ = true;
    *output_size = out - output;
    return true;
}

bool H264ToAnnexBBitstreamConverter::ConvertNalUnitStreamToByteStream(
    const uint8_t* input,
    uint32_t input_size,
    const mp4::AVCDecoderConfigurationRecord* avc_config,
    uint8_t* output,
    uint32_t* output_size)
{
    const uint8_t* inscan = input; // We read the input from here progressively
    uint8_t* outscan = output; // We write the output to here progressively
    uint32_t data_left = input_size;

    if (input_size == 0 || *output_size == 0) {
        *output_size = 0;
        return false; // Error: invalid input
    }

    // NAL unit width should be known at this point
    CHECK(nal_unit_length_field_width_ == 1 || nal_unit_length_field_width_ == 2 || nal_unit_length_field_width_ == 4);

    // Do the actual conversion for the actual input packet
    int nal_unit_count = 0;
    while (data_left > 0) {
        uint8_t i;
        uint32_t nal_unit_length;

        // Read the next NAL unit length from the input buffer by scanning
        // the input stream with the specific length field width
        for (nal_unit_length = 0, i = nal_unit_length_field_width_;
             i > 0 && data_left > 0;
             inscan++, i--, data_left--) {
            nal_unit_length <<= 8;
            nal_unit_length |= *inscan;
        }

        if (nal_unit_length == 0) {
            break; // Successful conversion, end of buffer
        } else if (nal_unit_length > data_left) {
            *output_size = 0;
            return false; // Error: not enough data for correct conversion
        }

        // Five least significant bits of first NAL unit byte signify
        // nal_unit_type.
        int nal_unit_type = *inscan & 0x1F;
        nal_unit_count++;

        // Insert the config after the AUD if an AUD is the first NAL unit or
        // before all NAL units if the first one isn't an AUD.
        if (avc_config && (nal_unit_type != H264NALU::kAUD || nal_unit_count > 1)) {
            uint32_t output_bytes_used = outscan - output;

            DCHECK_GE(*output_size, output_bytes_used);

            uint32_t config_size = *output_size - output_bytes_used;
            if (!ConvertAVCDecoderConfigToByteStream(*avc_config,
                    outscan,
                    &config_size)) {
                DVLOG(1) << "Failed to insert parameter sets.";
                *output_size = 0;
                return false; // Failed to convert the buffer.
            }
            outscan += config_size;
            avc_config = NULL;
        }
        uint32_t start_code_len;
        first_nal_unit_in_access_unit_ ? start_code_len = sizeof(kStartCodePrefix) + 1 : start_code_len = sizeof(kStartCodePrefix);
        if (static_cast<uint32_t>(outscan - output) + start_code_len + nal_unit_length > *output_size) {
            *output_size = 0;
            return false; // Error: too small output buffer
        }

        // Check if this packet marks access unit boundary by checking the
        // packet type.
        if (IsAccessUnitBoundaryNal(nal_unit_type)) {
            first_nal_unit_in_access_unit_ = true;
        }

        // Write extra zero-byte before start code prefix if this packet
        // signals next access unit.
        if (first_nal_unit_in_access_unit_) {
            *outscan = 0;
            outscan++;
            first_nal_unit_in_access_unit_ = false;
        }

        // No need to write leading zero bits.
        // Write start-code prefix.
        memcpy(outscan, kStartCodePrefix, sizeof(kStartCodePrefix));
        outscan += sizeof(kStartCodePrefix);
        // Then write the actual NAL unit from the input buffer.
        memcpy(outscan, inscan, nal_unit_length);
        inscan += nal_unit_length;
        data_left -= nal_unit_length;
        outscan += nal_unit_length;
        // No need for trailing zero bits.
    }
    // Successful conversion, output the freshly allocated bitstream buffer.
    *output_size = static_cast<uint32_t>(outscan - output);
    return true;
}

bool H264ToAnnexBBitstreamConverter::WriteParamSet(
    const std::vector<uint8_t>& param_set,
    uint8_t** out,
    uint32_t* out_size) const
{
    // Strip trailing null bytes.
    size_t size = param_set.size();
    while (size && param_set[size - 1] == 0)
        size--;
    if (!size)
        return false;

    // Verify space.
    uint32_t bytes_left = *out_size;
    if (bytes_left < kParamSetStartCodeSize || bytes_left - kParamSetStartCodeSize < size) {
        return false;
    }

    uint8_t* start = *out;
    uint8_t* buf = start;

    // Write the 4 byte Annex B start code.
    *buf++ = 0; // zero byte
    memcpy(buf, kStartCodePrefix, sizeof(kStartCodePrefix));
    buf += sizeof(kStartCodePrefix);

    // Copy the data.
    memcpy(buf, &param_set[0], size);
    buf += size;

    *out = buf;
    *out_size -= buf - start;
    return true;
}

} // namespace media
